# koroFileHeader at Yungoal acer
# Create: 2021-05-19 10:33:54
# LastEdit: 2021-05-20 10:52:19
"""使用 run command 的 REST API 接口，让虚拟机执行命令，获取命令结果
使用 运行命令 在虚拟机中运行脚本：
https://docs.microsoft.com/zh-cn/azure/virtual-machines/windows/run-command
REST API 接口：
https://docs.microsoft.com/zh-cn/rest/api/compute/virtual%20machines%20run%20commands/runcommand
获取 token 的方法：
https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-client-creds-grant-flow#get-a-token
"""
__author__ = '749B'

import os, sys
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 
sys.path.append(BASE_DIR)

from msrestazure.azure_cloud import AZURE_CHINA_CLOUD as CLOUD

import json
import requests
import time
from urllib.parse import urljoin

import pprint
# 加载环境变量
from environs import Env
env = Env()
env.read_env()

# 获取 token
token = '''
把 token 填在这里，避免重复获取
'''.strip()
token = None
if token is None:
    base_url = CLOUD.endpoints.active_directory
    url_path = '/{AZURE_TENANT_ID}/oauth2/v2.0/token'.format(AZURE_TENANT_ID=env('AZURE_TENANT_ID'))
    url = urljoin(base_url, url_path)
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    data = {
        'client_id': env('AZURE_CLIENT_ID'),
        'scope': CLOUD.endpoints.resource_manager + "/.default",
        'client_secret': env('AZURE_CLIENT_SECRET'),
        'grant_type': 'client_credentials',
    }
    r1 = requests.post(url, data=data, headers=headers)
    print(r1.json())
    token = r1.json().get('access_token')
assert token is not None

# 执行脚本
async_operation = '''
把 Azure-AsyncOperation 填在这里，避免重复获取
'''.strip()
async_operation = None
if async_operation is None:
    subscription_id = '064f5ddf-8a8f-4a99-a993-e2f88600f71d'
    resource_group_name = 'steed-dev'
    vm_name = 'winhost2'
    api_version = '2021-03-01'

    base_url = CLOUD.endpoints.resource_manager
    url_path = 'subscriptions/{subscription_id}/' \
        'resourceGroups/{resource_group_name}/' \
        'providers/Microsoft.Compute/virtualMachines/{vm_name}/' \
        'runCommand?api-version={api_version}'.format(
            subscription_id=subscription_id,
            resource_group_name=resource_group_name,
            vm_name=vm_name,
            api_version=api_version
        )
    url = urljoin(base_url, url_path)
    headers = {
        'Authorization': 'Bearer {token}'.format(token=token),
        'Content-Type': 'application/json',
    }
    script = """\
    param(
        [string]$arg1,
        [string]$arg2
    )
    Write-Host This is a sample script with parameters $arg1 and $arg2
    """
    data = {
        'commandId': 'RunPowerShellScript',
        'parameters': [{'name': 'arg1', 'value': 'v1'}, {'name': 'arg2', 'value': 'VALUE2'}],
        'script': script.splitlines(),
    }
    r2 = requests.post(
        url,
        data=json.dumps(data),
        headers=headers,
    )
    print(r2.status_code)
    if r2.text:
        print(r2.json())
    if r2.status_code == 202:
        # 这个返回为空，没有 r2.text
        print(r2.headers)
        async_operation = r2.headers.get('Azure-AsyncOperation')
    elif r2.status_code == 409:
        # 虚拟机未开机 {'error': {'code': 'OperationNotAllowed', 'message': 'The operation requires the VM to be running (or set to run).'}}
        # 还在执行命令 {'error': {'code': 'Conflict', 'message': 'Run command extension execution is in progress. Please wait for completion before invoking a run command.'}}
        pass
assert async_operation is not None

# 异步拿结果
headers = {
    'Authorization': 'Bearer {token}'.format(token=token),
    'Content-Type': 'application/json',
}
r3 = requests.get(async_operation, headers=headers)
retry = 1
while r3.ok and retry < 300:
    status = r3.json().get('status')
    if status.lower() == 'succeeded':
        pprint.pp(r3.json())
        break
    else:
        print(f'status: {status}. 等待 {retry} 秒...')
        time.sleep(retry)
        retry <<= 1
        r3 = requests.get(async_operation, headers=headers)
"""最后拿到的结果如下：
{'startTime': '2021-05-20T02:50:27.7159231+00:00',
 'endTime': '2021-05-20T02:50:58.0451754+00:00',
 'status': 'Succeeded',
 'properties': {'output': {'value': [{'code': 'ComponentStatus/StdOut/succeeded',
                                      'level': 'Info',
                                      'displayStatus': 'Provisioning succeeded',
                                      'message': 'This is a sample script with '
                                                 'parameters v1 and VALUE2'},
                                     {'code': 'ComponentStatus/StdErr/succeeded',
                                      'level': 'Info',
                                      'displayStatus': 'Provisioning succeeded',
                                      'message': ''}]}},
 'name': 'b1b46f11-a0f1-4a5b-af50-df19175d7503'}
"""
